#include <stdio.h>
#include <windows.h>
#include<iostream>
#include<stdlib.h>
using namespace std;

float img[10000][28][28];
int label[10000];

void qconv(int ch_in,int ch_out,int pad,int stride,int k,int h,int w,char* in,char* weight,int* bias,int* out){
    int h_o,w_o;
    int i,j,n,m;
    int kx,ky;
    h_o=(h-k+2*pad)/stride+1;
    w_o=(w-k+2*pad)/stride+1;
    for(i=0;i<h_o;i++)
        for(j=0;j<w_o;j++)
            for(m=0;m<ch_out;m++){
                //计算out[m][i][j]的地址
                int addr_o=m*h_o*w_o+i*w_o+j;
                *(out+addr_o)=bias[m];
                for(n=0;n<ch_in;n++)
                    for(kx=0;kx<k;kx++)
                        for(ky=0;ky<k;ky++)
                    {
                    //out[m][i][j]+=in[n][i*stride+kx-pad][j*stride+ky-pad]*weight[m][n][kx][ky]
                        int row=i*stride+kx-pad;
                        int col=j*stride+ky-pad;
                        int addr_i=n*h*w+row*w+col;
                        //
                        int addr_w=m*ch_in*k*k+n*k*k+kx*k+ky;
                        if(row>=0&&row<h&&col>=0&&col<w)
                            *(out+addr_o)+=(*(in+addr_i))*(*(weight+addr_w));
                    }
                *(out+addr_o)=(*(out+addr_o)>0.0)?(*(out+addr_o)):0.0;
    }
}
void conv(int ch_in,int ch_out,int pad,int stride,int k,int h,int w,float* in,float *weight,float *bias,float *out){
    int h_o,w_o;
    int i,j,n,m;
    int kx,ky;
    h_o=(h-k+2*pad)/stride+1;
    w_o=(w-k+2*pad)/stride+1;
    for(i=0;i<h_o;i++)
        for(j=0;j<w_o;j++)
            for(m=0;m<ch_out;m++){
                //计算out[m][i][j]的地址
                int addr_o=m*h_o*w_o+i*w_o+j;
                *(out+addr_o)=bias[m];
                for(n=0;n<ch_in;n++)
                    for(kx=0;kx<k;kx++)
                        for(ky=0;ky<k;ky++)
                    {
                    //out[m][i][j]+=in[n][i*stride+kx-pad][j*stride+ky-pad]*weight[m][n][kx][ky]
                        int row=i*stride+kx-pad;
                        int col=j*stride+ky-pad;
                        int addr_i=n*h*w+row*w+col;
                        //
                        int addr_w=m*ch_in*k*k+n*k*k+kx*k+ky;
                        if(row>=0&&row<h&&col>=0&&col<w)
                            *(out+addr_o)+=(*(in+addr_i))*(*(weight+addr_w));
                    }
                *(out+addr_o)=(*(out+addr_o)>0.0)?(*(out+addr_o)):0.0;
    }
}
void qpool(int h,int w,int ch,int* in,int* out){
    int i,j,n,kx,ky;
    for(i=0;i<h/2;i++)
        for(j=0;j<w/2;j++)
            for(n=0;n<ch;n++){
                //int tmp=in[n][2*i][2*j]
                int tmp1=*(in+n*h*w+2*i*w+2*j);
                int tmp2=*(in+n*h*w+2*i*w+2*j+1);
                int tmp3=*(in+n*h*w+(2*i+1)*w+2*j);
                int tmp4=*(in+n*h*w+(2*i+1)*w+2*j+1);
                int max1=(tmp1>tmp2)?tmp1:tmp2;
                int max2=(tmp3>tmp4)?tmp3:tmp4;
                *(out+n*h*w/4+i*w/2+j)=(max1>max2)?max1:max2;
    }
}
void pool(int h,int w,int ch,float* in,float* out){
    int i,j,n,kx,ky;
    for(i=0;i<h/2;i++)
        for(j=0;j<w/2;j++)
            for(n=0;n<ch;n++){
                //float tmp=in[n][2*i][2*j]
                float tmp1=*(in+n*h*w+2*i*w+2*j);
                float tmp2=*(in+n*h*w+2*i*w+2*j+1);
                float tmp3=*(in+n*h*w+(2*i+1)*w+2*j);
                float tmp4=*(in+n*h*w+(2*i+1)*w+2*j+1);
                float max1=(tmp1>tmp2)?tmp1:tmp2;
                float max2=(tmp3>tmp4)?tmp3:tmp4;
                *(out+n*h*w/4+i*w/2+j)=(max1>max2)?max1:max2;
    }
}
void global_avg_pool(float* in,int h,int w,int ch,float* out){
    for(int c=0;c<ch;c++){
        float tmp=0.0;
        for(int i=0;i<w;i++)
            for(int j=0;j<h;j++)
                tmp+=*(in+c*h*w+i*w+j);
        out[c]=tmp/(h*w);
    }
}
void read_param(float* buff,int length,const char* filename){
    FILE *fp;
    if((fp=fopen(filename,"rb"))==NULL){
        printf("Fail to open file!\n");
        exit(0);
    }
    fread(buff,sizeof(float),length,fp);
    fclose(fp);
}
void read_param_char(char* buff,int length,const char* filename){
    FILE *fp;
    float *tmp=new float[length];
    if((fp=fopen(filename,"rb"))==NULL){
        printf("Fail to open file!\n");
        exit(0);
    }
    fread(tmp,sizeof(float),length,fp);
    for(int i=0;i<length;i++)
        buff[i]=(char)tmp[i];
    delete [] tmp;
    fclose(fp);
}
void read_param_int(int* buff,int length,const char* filename){
    FILE *fp;
    float *tmp=new float[length];
    if((fp=fopen(filename,"rb"))==NULL){
        printf("Fail to open file!\n");
        exit(0);
    }
    fread(tmp,sizeof(float),length,fp);
    for(int i=0;i<length;i++)
        buff[i]=(int)tmp[i];
    fclose(fp);
}

void quant(float *in,float s,char *out,int length){
    for(int i=0;i<length;i++){
        int tmp=(int)(in[i]*s);
        if(tmp<-128)
            out[i]=(char)-128;
        else if(tmp>127)
            out[i]=(char)127;
        else
            out[i]=(char)tmp;
    }
}

void requantize(int* in,float* out,int ch,int h,int w,float As,float *Ws){
    for(int c=0;c<ch;c++)
        for(int i=0;i<h;i++)
            for(int j=0;j<w;j++)
                out[c*h*w+i*w+j]=(float)in[c*h*w+i*w+j]/(As*Ws[c]);
    return;
}

void requantize(int* in,char* out,int ch,int h,int w,float* s){      //s=Asn/(Asn-1*Wsn-1)
    for(int c=0;c<ch;c++)
        for(int i=0;i<h;i++)
            for(int j=0;j<w;j++)
                {
                    float tmp=(float)in[c*h*w+i*w+j]*s[c];
                    int qtmp=(int)tmp;
                    if(qtmp>127)
                        out[c*h*w+i*w+j]=(char)127;
                    else if(qtmp<-128)
                        out[c*h*w+i*w+j]=(char)-128;
                    else
                        out[c*h*w+i*w+j]=(char)qtmp;
                }
}

int main(){
    //parameters definition
    float *w0=new float[8*1*3*3];
    float *b0=new float[8];
    char* w1=new char[16*8*3*3];
    int* b1=new int[16];
    char* w2=new char[32*16*3*3];
    int* b2=new int[32];
    char* w3=new char[10*32*3*3];
    int* b3=new int[10];
    float As1,As2,As3;
    float Ws1[16];
    float Ws2[32];
    float Ws3[10];
    float s1[16];
    float s2[32];
    //parameters
    read_param(w0,8*1*3*3,"F:\\nn\\w0.bin");
    read_param(b0,8,"f:\\nn\\b0.bin");
    read_param_char(w1,16*8*3*3,"f:\\nn\\w1.bin");
    read_param_int(b1,16,"F:\\nn\\b1.bin");
    read_param_char(w2,32*16*3*3,"f:\\nn\\w2.bin");
    read_param_int(b2,32,"F:\\nn\\b2.bin");
    read_param_char(w3,10*32*3*3,"f:\\nn\\w3.bin");
    read_param_int(b3,10,"F:\\nn\\b3.bin");
    read_param(&As1,1,"F:\\nn\\As1.bin");
    read_param(&As2,1,"F:\\nn\\As2.bin");
    read_param(&As3,1,"F:\\nn\\As3.bin");
    read_param(Ws1,16,"F:\\nn\\Ws1.bin");
    read_param(Ws2,32,"F:\\nn\\Ws2.bin");
    read_param(Ws3,10,"F:\\nn\\Ws3.bin");
    //get s1 s2
    for(int i=0;i<16;i++)
        s1[i]=As2/(As1*Ws1[i]);
    for(int i=0;i<32;i++)
        s2[i]=As3/(As2*Ws2[i]);
    //data and label
    read_param((float*)img,10000*28*28,"f:\\nn\\img.bin");
    read_param_int(label,10000,"f:\\nn\\label.bin");
    //中间变量
    //conv0
    float *in=new float[28*28];                    
    float *conv0_out=new float[8*28*28];
    float *pool0_out=new float[8*14*14];
    //conv1
    char* conv1_in=new char[8*14*14];
    int* conv1_out=new int[16*14*14];
    float* requant_conv1_out=new float[16*14*14];
    //conv2
    char* conv2_in=new char[16*14*14];
    int* conv2_out=new int[32*14*14];
    float* requant_conv2_out=new float[32*14*14];
    //pool1
    int* pool1_out=new int[32*7*7];
    //conv3
    char* conv3_in=new char[32*7*7];
    int* conv3_out=new int[10*7*7];
    float* requant_conv3_out=new float[10*7*7];
    //global_avg_pool
    float* out=new float[10];
    //inference
    int cnt=0;
    for(int n=0;n<10000;n++){
        in=(float*)img[n];
        //conv0 and pool0
        conv(1,8,1,1,3,28,28,in,w0,b0,conv0_out);
        pool(28,28,8,conv0_out,pool0_out);
        //qconv1
        quant(pool0_out,As1,conv1_in,8*14*14);
        qconv(8,16,1,1,3,14,14,conv1_in,w1,b1,conv1_out);
        requantize(conv1_out,conv2_in,16,14,14,s1);
        //qconv2
        qconv(16,32,1,1,3,14,14,conv2_in,w2,b2,conv2_out);
        //pool1
        qpool(14,14,32,conv2_out,pool1_out);
        requantize(pool1_out,conv3_in,32,7,7,s2);
        //qconv3
        qconv(32,10,1,1,3,7,7,conv3_in,w3,b3,conv3_out);
        requantize(conv3_out,requant_conv3_out,10,7,7,As3,Ws3);
        //global avg pool
        global_avg_pool(requant_conv3_out,7,7,10,out);
        //argmax out
        float max_val=-9999;
        int max_idx=-1;
        for(int i=0;i<10;i++)
            if(out[i]>max_val)
           {
               max_val=out[i];
               max_idx=i;
           }
        //compare
        if(max_idx==label[n]){
            cnt++;
            printf("n=%d,right\n",n);
        }
    }
    cout<<"acc is "<<(float)cnt/10000.0<<endl;
    system("pause");
    return 0;
}
